import numpy as np
import pinecone as pc
from tensorflow.keras.datasets import mnist
from PIL import Image

# Pinecone 配置
PINECONE_API_KEY = "pcsk_2WNR5R_95XGmZAVXXQYSc2txi8pJwXVWQaUg2WmdnT1nREKVJRSXdTH4CaJ3Lq9NcfKwD3"
INDEX_NAME = "developer-quickstart-py"


def load_and_preprocess_sample():
    """加载并预处理一个 MNIST 测试样本（缩放为 8x8）"""
    (_, _), (test_images, _) = mnist.load_data()
    sample_img = test_images[0]  # 取第一个测试样本

    # 缩放并归一化
    img_pil = Image.fromarray(sample_img)
    img_resized = img_pil.resize((8, 8), Image.Resampling.LANCZOS)
    img_inverted = Image.eval(img_resized, lambda x: 255 - x)
    return np.array(img_inverted, dtype=np.float32) / 16  # 归一化到 0-16


if __name__ == "__main__":
    # 初始化 Pinecone 客户端
    pc_client = pc.Pinecone(api_key=PINECONE_API_KEY)
    index = pc_client.Index(INDEX_NAME)

    # 加载并预处理样本
    sample_vector = load_and_preprocess_sample().tolist()

    # 查询 Top-11 结果
    query_result = index.query(
        vector=sample_vector,
        top_k=11,
        include_metadata=True
    )

    # 打印结果并预测
    for match in query_result["matches"]:
        print(f"id: {match['id'].split('-')[-1]}, distance: {match['score']:.0f}, label: {match['metadata']['label']}")
    pred_labels = [int(match["metadata"]["label"]) for match in query_result["matches"]]
    predicted_digit = max(set(pred_labels), key=pred_labels.count)
    print(f"Predicted digit: {predicted_digit}")